package cn.lsh.spark.sql;

import org.apache.spark.sql.Row;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructType;

import java.util.Arrays;

public class CountStrLen extends UserDefinedAggregateFunction {

	private static final long serialVersionUID = -2117767025657696538L;

	@Override
	public StructType inputSchema() {
		//指定输入字段及类型
		return DataTypes.createStructType(Arrays.asList(DataTypes.createStructField("field", DataTypes.StringType, true)));
	}

	@Override
	public StructType bufferSchema() {
		//指定聚合过程中MutableAggregationBuffer的字段及类型
		return DataTypes.createStructType(Arrays.asList(DataTypes.createStructField("field1", DataTypes.IntegerType, true)));
	}

	@Override
	public DataType dataType() {
		//指定UDAF函数计算后返回的结果类型
		return DataTypes.IntegerType;
	}

	@Override
	public boolean deterministic() {
		//确保一致性一般用true，用以标记针对给定的一组输入，UDAF是否总是生成相同的结果
		return true;
	}

	@Override
	public void initialize(MutableAggregationBuffer buffer) {
		//初始化一个自己内部定义的值，在聚合之前每组数据初始化结果
		buffer.update(0,0);
	}

	@Override
	public void update(MutableAggregationBuffer buffer, Row input) {
		//更新，可以认为一个一个地将组内的字段值传递进来，实现拼接的逻辑
		//buffer.getInt(0)获得上一组聚合后的结果，相当于map端的combiner
		//这里即是：在进行聚合的时候，每当有新的值进来，对分组后的聚合如何进行计算
		//buffer是上一轮输入数据的聚合输出，input是方法当前输入的数据
		String field = input.getString(0);
		int len = field != null ? field.length() : 0;
		buffer.update(0, buffer.getInt(0) + len);
	}

	@Override
	public void merge(MutableAggregationBuffer buffer, Row row) {
		//合并上面的update操作，buffer是一组数据update的输出，row是另一组数据的update输出
		buffer.update(0, buffer.getInt(0) + row.getInt(0));
	}

	@Override
	public Object evaluate(Row row) {
		//最后返回一个和datatype方法的类型要一致的类型，返回UDAF最后的计算结果
		//入参row是merge输出的数据
		return row.get(0);
	}
}
